#include <stdio.h>
#include <stdlib.h>
#include "util.h"
#include "rdiff.h"
#include "common.h"
#include "sum_sum2_kernel.h"
#include "corrcoefmatrix_kernel.h"

#define ROWSIZE 1024

/* DEBUG will print the entire output matrix */
//#define DEBUG 

float threshold = 0.1;

void usage(int argc, char *argv[]) {
   fprintf(stderr,"Usage: %s <casefile> <ctrlfile> <numrow> <numcol> <threshold>\n",argv[0]);
}

/*
 * Reserve and create sum and sum^2 in device memory
 */
void create_sum_sum2(char *filename,int nrow, int ncol,float **sumG,float **sum2G) {
   snps_data_t  *snps_data;
   int padded_nrow = ((nrow + ROWSIZE -1)/ROWSIZE)*ROWSIZE;
   double st,et;

   cudaError_t err;
   st = getTimeStamp();
   fprintf(stderr,"Reserve device memory for sum and sum2: %.2f MB.\n",2*padded_nrow*sizeof(float)/1E6);
   if ((err=cudaMalloc((void **) sumG,padded_nrow*sizeof(float))) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }
   if ((err=cudaMalloc((void **) sum2G,padded_nrow*sizeof(float))) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }

   dim3 blockDim(TILESIZE,TILESIZE);
   dim3 gridDim(1,ROWSIZE/TILESIZE);

   snps_data = new_snps(filename,ncol);
   int blksize = ROWSIZE*ncol*sizeof(float);
   float *blk = (float *) malloc(blksize); 
   float *blkG;

   if ((err=cudaMalloc((void **) &blkG,blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }
   et = getTimeStamp();
   fprintf(stderr,"Reserve memory spends time %.2f ms.\n",et-st);

   st = getTimeStamp();
   int n = 0;
   int offset = 0;
   while ((n = read_next_snps(snps_data,ROWSIZE,blk,blksize)) > 0) {
      cudaMemcpy(blkG,blk,blksize,cudaMemcpyHostToDevice);
      sum_sum2_kernel<<<gridDim,blockDim>>>(*sumG+offset,*sum2G+offset,blkG,n,ncol);     
      offset += ROWSIZE;
   }
   cudaThreadSynchronize();
   et = getTimeStamp();
   fprintf(stderr,"Calculate sum and sum2 for %s spends time %.2f ms.\n",snps_data->filename,et-st);

   // check for error
   check_cuda_errors(__FILE__,__LINE__);
   cudaFree(blkG);
   free(blk);
   close_snps(snps_data);
   free(snps_data);

}
/*
 * Do processing on a block of output page
 * @param nrow Number of rows in the entire output (input)
 * @param ncol Number of columns in the entire output (input)
 * @param blk The output page (input)
 * @param blkrow Number of rows in output page (input)
 * @param blkcol Number of columns in output page (input)
 * @param ll_i The offset of lower left corner of the output page in y-axis (row) (input)
 * @param ll_j the offset of lower left corner of the output page in x-axis (column) (input)
 */
void do_processing(int nrow, int ncol, float *blk, int blkrow, int blkcol, int ll_i, int ll_j) {
    int i,j;
    int gi,gj; // global row and column index
    float r;

    for (i=0; i < blkrow; i++) {
       gi = (ll_i + i);
       for (j=0; j < blkcol; j++) {
          gj = (ll_j + j);
          if ( (gi < nrow) && (gj < ncol) ) {
             if (gi >= gj) {
//                blk[i*blkcol + j] = 0.0; 
             } else {
                 r = blk[i*blkcol+j];
                 if (r >= threshold) {
                    printf("%d %d %f\n",gi,gj,r);
                 }
             }
          }
       }
    }

}

void put_submatrix(float *A,int nrow, int ncol, float *blk, int blkrow, int blkcol, int ll_i, int ll_j) {
    int i,j;
    int gi,gj; // global row and column index
    for (i=0; i < blkrow; i++) {
       gi = (ll_i + i);
       for (j=0; j < blkcol; j++) {
          gj = (ll_j + j);
          if ((gi < nrow) && (gj < ncol)) {
             A[gi*ncol + gj] = blk[i*blkcol + j]; 
          }
       }
    }
}

int main(int argc, char *argv[]) {
   char *ctrl_filename, *case_filename;
   int ncol,nrow;
   float threshold;
   int size,blksize;
   int i,j;
   snps_data_t *ctrl_snps_i,*ctrl_snps_j;
   snps_data_t *case_snps_i,*case_snps_j;
   float *ctrlblk_i,*ctrlblk_j;
   float *caseblk_i,*caseblk_j;
   double st,et;

   if ((ROWSIZE%TILESIZE) > 0) {
      fprintf(stderr,"ROWSIZE must be a multiple of TILESIZE\n");
      exit(1);
   }

   if (argc < 6) {
      usage(argc,argv);
      exit(1);
   }

   case_filename = argv[1];
   ctrl_filename = argv[2];
   nrow = atoi(argv[3]); 
   ncol = atoi(argv[4]);
   threshold = (float) atof(argv[5]);
   ctrl_snps_i = new_snps(ctrl_filename,ncol);
   ctrl_snps_j = new_snps(ctrl_filename,ncol);
   case_snps_i = new_snps(case_filename,ncol);
   case_snps_j = new_snps(case_filename,ncol);

   fprintf(stderr,"PAGE ROWSIZE = %d, THREAD BLOCK SIZE = %d\n",ROWSIZE, TILESIZE);
   fprintf(stderr,"Threshold = %f\n",threshold);

   float *ctrl_sumG, *ctrl_sum2G;
   create_sum_sum2(ctrl_filename,nrow,ncol,&ctrl_sumG,&ctrl_sum2G);
/*
   size = nrow*sizeof(float);
   float *ctrl_sum = (float *) malloc(size); 
   float *ctrl_sum2 = (float *) malloc(size);
   cudaMemcpy(ctrl_sum,ctrl_sumG,size,cudaMemcpyDeviceToHost);
   cudaMemcpy(ctrl_sum2,ctrl_sum2G,size,cudaMemcpyDeviceToHost);
*/

   float *case_sumG, *case_sum2G;
   create_sum_sum2(case_filename,nrow,ncol,&case_sumG,&case_sum2G);

   st = getTimeStamp();
   blksize = ncol*ROWSIZE*sizeof(float);
   ctrlblk_i = (float *) malloc(blksize);
   ctrlblk_j = (float *) malloc(blksize);
   caseblk_i = (float *) malloc(blksize);
   caseblk_j = (float *) malloc(blksize);

   cudaError_t err;
   float *ctrlblk_iG, *ctrlblk_jG;
   fprintf(stderr,"Reserve device memory for ctrlblk i and j: %.2f MB.\n",2*blksize/1E6);
   if ((err=cudaMalloc((void **) &ctrlblk_iG,blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }
   if ((err=cudaMalloc((void **) &ctrlblk_jG,blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }

   float *caseblk_iG, *caseblk_jG;
   fprintf(stderr,"Reserve device memory for caseblk i and j: %.2f MB.\n",2*blksize/1E6);
   if ((err=cudaMalloc((void **) &caseblk_iG,blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }
   if ((err=cudaMalloc((void **) &caseblk_jG,blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }

   int output_blksize = ROWSIZE*ROWSIZE*sizeof(float);
   float *ctrlblk_coeff,*ctrlblk_coeffG;
   ctrlblk_coeff = (float *) malloc(output_blksize);
   fprintf(stderr,"Reserve device memory for correlation matrix for ctrlblk: %.2f MB.\n",output_blksize/1E6);
   if ((err=cudaMalloc((void **) &ctrlblk_coeffG,output_blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }

   float *caseblk_coeff,*caseblk_coeffG;
   caseblk_coeff = (float *) malloc(output_blksize);
   fprintf(stderr,"Reserve device memory for correlation matrix for caseblk: %.2f MB.\n",output_blksize/1E6);
   if ((err=cudaMalloc((void **) &caseblk_coeffG,output_blksize)) != cudaSuccess) {
      fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
      exit(1);
   }
   et = getTimeStamp();
   fprintf(stderr,"Reserve memory spends %.2f ms.\n",et-st);

   int grid_size_x = ROWSIZE/TILESIZE;
   int grid_size_y = ROWSIZE/TILESIZE;
   dim3 blockDim(TILESIZE,TILESIZE);
   dim3 gridDim(grid_size_x,grid_size_y);
   fprintf(stderr,"Set up grid size %dx%d for each page\n",grid_size_x,grid_size_y);

   int offset_i=0, offset_j=0;

#ifdef DEBUG
   float *A = (float *) malloc(nrow*nrow*sizeof(float));
   fprintf(stderr,"Reserve host memory for debuging: %.2f MB\n",nrow*nrow*sizeof(float)/1E6);
#endif  
 
   int page_i = 0;
   int page_j = 0;
   
   st = getTimeStamp();
   int ctrl_n_i, ctrl_n_j;
   int case_n_i, case_n_j;
   double st1,et1;
   st1 = getTimeStamp();
   while ( ((ctrl_n_i = read_next_snps(ctrl_snps_i,ROWSIZE,ctrlblk_i,blksize)) > 0) && ((case_n_i = read_next_snps(case_snps_i,ROWSIZE,caseblk_i,blksize)) > 0) ) {
      cudaMemcpy(ctrlblk_iG,ctrlblk_i,blksize,cudaMemcpyHostToDevice);
      cudaMemcpy(caseblk_iG,caseblk_i,blksize,cudaMemcpyHostToDevice);
      et1 = getTimeStamp();
      fprintf(stderr,"Time to read block i %d = %.2f ms.\n",offset_i,et1-st1);
      st1 = getTimeStamp();
      while ( ((ctrl_n_j = read_next_snps(ctrl_snps_j,ROWSIZE,ctrlblk_j,blksize)) > 0) && ((case_n_j = read_next_snps(case_snps_j,ROWSIZE,caseblk_j,blksize)) > 0) ) {
         et1 = getTimeStamp();
         fprintf(stderr,"Time to read block j %d = %.2f ms.\n",offset_j,et1-st1);
         if (offset_i < offset_j+ROWSIZE-1) { 
            // compute only block overlapped with upper triangle
            st1 = getTimeStamp(); 
            cudaMemcpy(ctrlblk_jG,ctrlblk_j,blksize,cudaMemcpyHostToDevice);
            cudaMemcpy(caseblk_jG,caseblk_j,blksize,cudaMemcpyHostToDevice);
            corrcoefmatrixAB_kernel<<<gridDim,blockDim>>>(ctrlblk_coeffG,ctrlblk_iG,ctrlblk_jG,ctrl_sumG+offset_i,ctrl_sum2G+offset_i,ctrl_sumG+offset_j,ctrl_sum2G+offset_j,ROWSIZE,ncol);
            corrcoefmatrixAB_kernel<<<gridDim,blockDim>>>(caseblk_coeffG,caseblk_iG,caseblk_jG,case_sumG+offset_i,case_sum2G+offset_i,case_sumG+offset_j,case_sum2G+offset_j,ROWSIZE,ncol);

//            cudaThreadSynchronize();  // do we need this????
            sub_kernel<<<(ROWSIZE*ROWSIZE+256-1)/256,256>>>(caseblk_coeffG,ctrlblk_coeffG,ROWSIZE*ROWSIZE);
//            cudaMemcpy(ctrlblk_coeff,ctrlblk_coeffG,output_blksize,cudaMemcpyDeviceToHost);
           
            cudaMemcpy(caseblk_coeff,caseblk_coeffG,output_blksize,cudaMemcpyDeviceToHost);

            et1 = getTimeStamp();
            fprintf(stderr,"Time to calculate block %d,%d = %.2f ms.\n",offset_i,offset_j,et1-st1);
            st1 = getTimeStamp();
            do_processing(nrow,ncol,caseblk_coeff,ROWSIZE,ROWSIZE,offset_i,offset_j);
            et1 = getTimeStamp();
            fprintf(stderr,"Time to do_processing %.2f ms.\n",et1-st1);

         } else {
#ifdef DEBUG
            memset(caseblk_coeff,0,ROWSIZE*ROWSIZE*sizeof(float));
#endif
         }
 
#ifdef DEBUG 
         put_submatrix(A,nrow,nrow,caseblk_coeff,ROWSIZE,ROWSIZE,offset_i,offset_j);
#endif

         offset_j += ROWSIZE;
         page_j++;
      }
      if (ctrl_n_j < 0) {
         fprintf(stderr,"Fail to read file %s\n",ctrl_snps_j->filename);
         exit(1);
      }
      if (case_n_j < 0) {
         fprintf(stderr,"Fail to read file %s\n",case_snps_j->filename);
         exit(1);
      }
      reset_snps(ctrl_snps_j);
      reset_snps(case_snps_j);
      offset_j = 0;
      page_j = 0;
      offset_i += ROWSIZE;
      page_i++;
   }
   cudaThreadSynchronize();
   et = getTimeStamp();
   fprintf(stderr,"GPU calucation spends %.2f ms.\n",et-st);

   if (ctrl_n_i < 0) {
      fprintf(stderr,"Fail to read file %s\n",ctrl_snps_i->filename);
      exit(1);
   }
   if (case_n_i < 0) {
      fprintf(stderr,"Fail to read file %s\n",case_snps_i->filename);
      exit(1);
   }

   cudaFree(ctrlblk_iG);
   cudaFree(ctrlblk_jG);
   cudaFree(ctrlblk_coeffG);
   cudaFree(caseblk_iG);
   cudaFree(caseblk_jG);
   cudaFree(caseblk_coeffG);
   cudaFree(ctrl_sumG);
   cudaFree(ctrl_sum2G);
   cudaFree(case_sumG);
   cudaFree(case_sum2G);
   free(ctrlblk_coeff);
   free(caseblk_coeff);
   free(ctrlblk_i);
   free(ctrlblk_j);
   free(caseblk_i);
   free(caseblk_j);

#ifdef DEBUG
   for (i=0; i < nrow; i++) {
      for (j=0; j < nrow; j++) {
         printf("%f ",A[i*nrow+j]);
      }
      printf("\n");
   }
#endif

   fprintf(stderr,"%d total lines read, %d rows/page\n",ctrl_snps_i->current_row+1,page_i);

}
